#!/usr/env/bin python
"""
:mod:`pynchon.core.symbol`
==========================

Provides array objects to store N-dimensional string symbols.
"""

import numpy as np
from itertools import repeat
from pynchon.plot.shell import print_strings
from entity import Entity_maker, Entity_sequencer, Entity_lohi


class Symbol(object):
    """
    Symbol interface class. This class provides methods specific to 
    N-dimensional symbols e.g. 0D - 'A', 1D - ('A',), 2D - ('A', 'A').
    
    Provides: 'raw' and 'is_valid' methods and the 'rank' attribute to
    subclasses. 
    """
    rank = 0

    @classmethod
    def is_valid_element(cls, element, metadata):
        """
        Verify if element can be a valid ``Symbol`` instance.        
        """
        if not np.iterable(element) or isinstance(element, basestring):
            element = tuple([letter for letter in element])
        for symbol in element:
            if not isinstance(symbol, basestring) or len(symbol) > 1:
                raise ValueError(\
                "Not all elements of \"%s\" are a symbols." % (element,))

    @classmethod
    def is_valid_array(cls, array, metadata):
        alphabet = metadata.get('alphabet', None)
        if alphabet is not None:
            # this loop is very expensive (TODO)
            for element in array.flat:
                if not np.iterable(element):
                    element = (element,)
                for index in element:
                    if not index in alphabet:
                        raise ValueError(\
                       'Not all elements of "%s" are in alphabet.' % (element,))
        if (np.issubdtype(array.dtype, 'S') and \
            (array.dtype.itemsize != (len(array.dtype) or 1))):
            raise ValueError('Only "S1" and "U1" type symbols are supported.')
        if (np.issubdtype(array.dtype, 'U') and \
            (array.dtype.itemsize != 4*(len(array.dtype) or 1))):
            raise ValueError('Only "S1" and "U1" type symbols are supported.')
        
    @classmethod
    def raw_sequence(cls, rank, dtype, fval, metadata, fast, kwargs, entities):
        cls.is_valid_element(entities, metadata)
        a_dims = len(dtype)     # array dimensions
        e_dims = len(entities)  # entity dimensions
        try:
            if ((a_dims and (a_dims != e_dims)) or \
               (not a_dims and (e_dims > 1))):
                    # it's complex because if len(dtype) == 0 then len(entites)
                    # should be 1 or len(dtype) should be len(entities) if 
                    # len(dtype) is >= 1.
                    raise ValueError
            if not "".join(entities):
                # np.array('',dtype='S1') vs np.array('',dtype='c')
                entities = np.array(entities, dtype='S1').view(dtype)
            else:
                entities = np.array(entities, dtype='c').view(dtype)
            entities.shape = ()
            #entities = np.array(entities, dtype=dtype)
        except (ValueError, TypeError):
            raise ValueError("Expected a %sD sequence of symbols got: \"%s\"." \
                              % (a_dims, entities))
        cls.is_valid_array(entities, metadata)
        cls.is_valid_array(fval, metadata)
        return (entities,)

    def __str__(self):
        return "".join(self.item())

    def recode(self, code, inplace=False):
        """
        This re-codes (translates) the symbol given a "code" dictionary. Returns
        a new instance if "inplace" is ``False``. 
        
        Arguments:
        
            - code (``dict``) Translation dictionary.
            - inplace (``bool``) If ``True`` the current instance will have the 
              new value
        """
        if inplace:
            symbol = code[self.item()]
            self[None][0] = symbol
        else:
            cpy = self.copy()
            cpy.recode(code, inplace =True)
            return cpy

    def _make_elements(self, elements):
        # TODO make this uniform
        if isinstance(elements, tuple):
            elements = list(elements)
        try:
            elements = np.asarray(elements, dtype=self.dtype)
        except ValueError:
            raise ValueError('Elements are not compatible with the dtype: "%s"' % self.dtype)
        return elements

    def _keep_or_omit(self, elements, mode):
        elements = self._make_elements(elements)
        inelements = (self == elements).any()
        
        if (inelements and (mode == 'omit')) or \
           (not inelements and (mode == 'keep')):
            empty = tuple(repeat('',len(self.dtype))) if len(self.dtype) else ""
            return np.array(empty, dtype=self.dtype)
        else:
            return self
        
    def keep(self, elements):
        """
        Returns self (not a copy) if in "elements" else returns an empty 
        ``Symbol`` instance of the correct data type. 
        
        Arguments:
        
            - elements (sequence or any ``Symbol`` instance) A numpy array of 
              the same (or compatible) data type as symbol or a sequence of 
              elements, which can be made into one."""
        return self._keep_or_omit(elements, 'keep')
        
    def omit(self, elements):
        """
        Returns self (not a copy) if not in "elements" else returns an empty 
        ``Symbol`` instance of the correct data type. 
        
        Arguments:
        
            - elements (sequence or any ``Symbol`` instance) A numpy array of 
              the same (or compatible) data type as symbol or a sequence of 
              elements, which can be made into one.
        """
        return self._keep_or_omit(elements, 'omit')


class SymbolVector(Symbol):
    """
    SymbolVector interface class. This class provides methods specific to
    vectors of N-dimensional symbols e.g. 0D - ['A', 'B', ... , 'Z'],
    1D - [('A',), ('B',), ... ('Z',)], 2D - [('A', 'a'), ('B', 'b'), ... ,
    ('Z', 'z')].
    """
    rank = 1

    @classmethod
    def raw_sequence(cls, rank, dtype, fval, metadata, fast, kwargs, entities):
        a_dims = len(dtype) or 1 # for xrange step cannot be 0
        if isinstance(entities, basestring):
            # got string
            # slice it
            entities_generator = (entities[idx:idx+a_dims] for idx in xrange(0,len(entities), a_dims))
            if fast:
                output = [np.array(sub, dtype=dtype) for sub in entities_generator]
            else:
                output = Entity_sequencer(cls, dtype, fval, metadata, entities_generator)
        else:
            if fast:
                # if fast make self from arrays
                output = entities
            else:
                # if slow make arrays into Symbols and self from Symbols 
                output = Entity_sequencer(cls, dtype, fval, metadata, entities)
        return output
           
    def _keep_or_omit(self, elements, mode):
        indices = []
        lo_clss, hi_clss = Entity_lohi(self)
        lo = lo_clss[0]
        elements = self._make_elements(elements)
        empty = np.array(tuple(repeat('',len(self.dtype))) if len(self.dtype) else "", \
                         dtype =self.dtype)
        for idx in xrange(len(self)):
            symbol = self[idx:idx+1].view(lo).reshape(())
            if symbol._keep_or_omit(elements, mode) != empty:
                indices.append(idx)
        if indices:
            return self[np.array(indices)]
        else:
            return self[0:0]

    def recode(self, code, inplace=False):
        if inplace:
            # np ravel might return a copy
            flat = np.ravel(self)
            if not np.may_share_memory(self, flat):
                raise ValueError("This array cannot be recoded in place.")
        else:
            flat = np.ravel(self.copy())
        for idx, symbol in enumerate(flat):
            flat[idx] = code[symbol.item()]
        if not inplace:
            return flat.reshape(self.shape)
        
    def keep(self, elements):
        """
        Returns a view only of elements which are in "elements".
        
        Arguments:
        
            - elements see: ``Symbol.keep``
        """        
        return self._keep_or_omit(elements, 'keep')    
                               
    def omit(self, elements):
        """
        Returns a view with elements which are in "elements" omitted.
        
        Arguments:
        
            - elements see: ``Symbol.omit``
        """
        return self._keep_or_omit(elements, 'omit')

    def find(self, ktup):
        """
        Determines the start positions of all occurrences of ktup in ``self``.
        The "ktup" argument should be a sequence of symbols of compatible 
        data type. Returns a sequence of arrays like the output of 
        ``numpy.where``.
        
        Arguments:
        
            - ktup (sequence of symbols) see: above and 
              ``Symbol._make_elements``
        """
        # first index from right which cannot True 
        end = -len(ktup) + 1
        # the output array is of the same shape as the input.
        starts = np.ndarray(shape=self.shape, dtype='bool')
        # all should be True
        starts.fill(True)
        # only the (k-1)-last elements cannot be.
        if end: 
            # len(ktup) > 1 
            starts[...,end:].fill(False)
            spart = starts[...,:end]
        else:
            spart = starts
        for i, element in enumerate(ktup):
            bools = (self == element)[...,i or None:end+i or None]
            np.logical_and(spart, bools, spart) # this replaces the results in s
        return np.where(starts)

class SymbolArray(SymbolVector):
    """
    SymbolArray interface class. This class provides methods specific to
    arrays of N-dimensional symbols.
    
    """
    rank = 2
    
    def omit(self, elements):
        """
        this will need to be a copy
        """
        pass


class SymbolCube(SymbolArray):
    """
    SymbolCube interface class. This class provides methods specific to
    cubes (3D) of N-dimensional symbols.
    """
    rank = 3


class SymbolHyperCube(SymbolCube):
    """
    SymbolHyperCube interface class. This class provides methods specific to
    cubes (4D) of N-dimensional symbols.
    """
    rank = 4


# create classes
classes = Entity_maker(base_classes=(Symbol, SymbolVector, SymbolArray, \
                                       SymbolCube, SymbolHyperCube),
                        dims=(0, 1, 2, 3, 4),
                        types=('S1', 'U1')
                        # Similarly, creating a PyArray_CHAR array from the C 
                        # API results in dtype("|S1").
                       )
# add to globals
for cls in classes:
    globals()[cls.__name__] = cls




















